"""
XKCD plot generator
-------------------
Author: Jake Vanderplas

This is a script that will take any matplotlib line diagram, and convert it
to an XKCD-style plot.  It will work for plots with line & text elements,
including axes labels and titles (but not axes tick labels).

The idea for this comes from work by Damon McDougall
  http://www.mail-archive.com/matplotlib-users@lists.sourceforge.net/msg25499.html
"""
import numpy as np
import pylab as pl
from scipy import interpolate, signal
import matplotlib.font_manager as fm
from matplotlib import patches, transforms, text, patheffects
from matplotlib.artist import getp
from random import random

# We need a special font for the code below.  It can be downloaded this way:
import os
import urllib2
#if not os.path.exists('Humor-Sans.ttf'):
    #fhandle = urllib2.urlopen('http://antiyawn.com/uploads/Humor-Sans.ttf')
#    open('Humor-Sans.ttf', 'wb').write(fhandle.read())

    
def xkcd_line(x, y, xlim=None, ylim=None,
              mag=1.0, f1=30, f2=0.05, f3=15):
    """
    Mimic a hand-drawn line from (x, y) data

    Parameters
    ----------
    x, y : array_like
        arrays to be modified
    xlim, ylim : data range
        the assumed plot range for the modification.  If not specified,
        they will be guessed from the  data
    mag : float
        magnitude of distortions
    f1, f2, f3 : int, float, int
        filtering parameters.  f1 gives the size of the window, f2 gives
        the high-frequency cutoff, f3 gives the size of the filter
    
    Returns
    -------
    x, y : ndarrays
        The modified lines
    """
    x = np.asarray(x)
    y = np.asarray(y)
    
    # get limits for rescaling
    if xlim is None:
        xlim = (x.min(), x.max())
    if ylim is None:
        ylim = (y.min(), y.max())

    if xlim[1] == xlim[0]:
        xlim = ylim
        
    if ylim[1] == ylim[0]:
        ylim = xlim

    # scale the data
    x_scaled = (x - xlim[0]) * 1. / (xlim[1] - xlim[0])
    y_scaled = (y - ylim[0]) * 1. / (ylim[1] - ylim[0])

    # compute the total distance along the path
    dx = x_scaled[1:] - x_scaled[:-1]
    dy = y_scaled[1:] - y_scaled[:-1]
    dist_tot = np.sum(np.sqrt(dx * dx + dy * dy))

    # number of interpolated points is proportional to the distance
    Nu = int(200 * dist_tot)
    u = np.arange(-1, Nu + 1) * 1. / (Nu - 1)

    # interpolate curve at sampled points
    k = min(3, len(x) - 1)
    res = interpolate.splprep([x_scaled, y_scaled], s=0, k=k)
    x_int, y_int = interpolate.splev(u, res[0]) 

    # we'll perturb perpendicular to the drawn line
    dx = x_int[2:] - x_int[:-2]
    dy = y_int[2:] - y_int[:-2]
    dist = np.sqrt(dx * dx + dy * dy)

    # create a filtered perturbation
    coeffs = mag * np.random.normal(0, 0.01, len(x_int) - 2)
    b = signal.firwin(f1, f2 * dist_tot, window=('kaiser', f3))
    response = signal.lfilter(b, 1, coeffs)

    x_int[1:-1] += response * dy / dist
    y_int[1:-1] += response * dx / dist

    # un-scale data
    x_int = x_int[1:-1] * (xlim[1] - xlim[0]) + xlim[0]
    y_int = y_int[1:-1] * (ylim[1] - ylim[0]) + ylim[0]
    
    return x_int, y_int

def DaVinci_line(line, color = "k", xlim = None, ylim = None):
    mag=1.0
    f1=50
    f2=0.01
    f3=15
    
    lines = [[],[]]
    
    x, y = line.get_data()

    lw = line.get_linewidth()
    
    ## background line for axes
    for i in xrange(20):
        x_int, y_int = xkcd_line(x, y, xlim, ylim, 
                             mag*1.5, f1, f2, f3)
        line_bg = pl.Line2D(x_int, y_int, color="#A68358",
                        linewidth=1.2*lw, alpha = 0.5)
        lines[1].append(line_bg)
    
   
    
    for i in xrange(4):
        x_int, y_int = xkcd_line(x, y, xlim, ylim, 
                             mag/2, f1, f2, f3)
        line_fg = pl.Line2D(x_int, y_int, color=line.get_color(),
                        linewidth= 2*lw, alpha = 0.6, linestyle = line.get_linestyle())
        lines[0].append(line_fg)
        
     ## Foreground lines
    x_int, y_int = xkcd_line(x, y, xlim, ylim, 
                             mag/2, f1, f2, f3)
    line_fg = pl.Line2D(x_int, y_int, color='k',
                    linewidth= 1*lw, alpha = 0.8, linestyle = line.get_linestyle())
    lines[0].append(line_fg)
    
    
    return lines
    
def XKCDify(ax, mag=1.0,
            f1=50, f2=0.01, f3=15,
            bgcolor='w',
            xaxis_loc=None,
            yaxis_loc=None,
            xaxis_arrow='+',
            yaxis_arrow='+',
            ax_extend=0.1,
            expand_axes=False):
    """Make axis look hand-drawn

    This adjusts all lines, text, legends, and axes in the figure to look
    like xkcd plots.  Other plot elements are not modified.
    
    Parameters
    ----------
    ax : Axes instance
        the axes to be modified.
    mag : float
        the magnitude of the distortion
    f1, f2, f3 : int, float, int
        filtering parameters.  f1 gives the size of the window, f2 gives
        the high-frequency cutoff, f3 gives the size of the filter
    xaxis_loc, yaxis_log : float
        The locations to draw the x and y axes.  If not specified, they
        will be drawn from the bottom left of the plot
    xaxis_arrow, yaxis_arrow : str
        where to draw arrows on the x/y axes.  Options are '+', '-', '+-', or ''
    ax_extend : float
        How far (fractionally) to extend the drawn axes beyond the original
        axes limits
    expand_axes : bool
        if True, then expand axes to fill the figure (useful if there is only
        a single axes in the figure)
    """
    # Get axes aspect
    ext = ax.get_window_extent().extents
    aspect = (ext[3] - ext[1]) / (ext[2] - ext[0])

    xlim = ax.get_xlim()
    ylim = ax.get_ylim()

    xspan = xlim[1] - xlim[0]
    yspan = ylim[1] - xlim[0]

    xax_lim = (xlim[0] - ax_extend * xspan,
               xlim[1] + ax_extend * xspan)
    yax_lim = (ylim[0] - ax_extend * yspan,
               ylim[1] + ax_extend * yspan)

    if xaxis_loc is None:
        xaxis_loc = ylim[0]

    if yaxis_loc is None:
        yaxis_loc = xlim[0]

    # Draw axes
    xaxis = pl.Line2D([xax_lim[0], xax_lim[1]], [xaxis_loc, xaxis_loc],
                      linestyle='-', color='k')
    yaxis = pl.Line2D([yaxis_loc, yaxis_loc], [yax_lim[0], yax_lim[1]],
                      linestyle='-', color='k')

    # Label axes3, 0.5, 'hello', fontsize=14)
    ax.text(xax_lim[1], xaxis_loc - 0.02 * yspan, ax.get_xlabel(),
            fontsize=14, ha='right', va='top', rotation=12)
    ax.text(yaxis_loc - 0.02 * xspan, yax_lim[1], ax.get_ylabel(),
            fontsize=14, ha='right', va='top', rotation=78)
    ax.set_xlabel('')
    ax.set_ylabel('')

    # Add title
    ax.text(0.5 * (xax_lim[1] + xax_lim[0]), yax_lim[1],
            ax.get_title(),
            ha='center', va='bottom', fontsize=16)
    ax.set_title('')

    Nlines = len(ax.lines)
    lines = [xaxis, yaxis] + [ax.lines.pop(0) for i in range(Nlines)]

    for line in lines:
        x, y = line.get_data()

        x_int, y_int = xkcd_line(x, y, xlim, ylim,
                                 mag, f1, f2, f3)

        # create foreground and background line
        lw = line.get_linewidth()
        line.set_linewidth(2 * lw)
        line.set_data(x_int, y_int)

        # don't add background line for axes
        if (line is not xaxis) and (line is not yaxis):
            line_bg = pl.Line2D(x_int, y_int, color=bgcolor,
                                linewidth=8 * lw)

            ax.add_line(line_bg)
        ax.add_line(line)

    # Draw arrow-heads at the end of axes lines
    arr1 = 0.03 * np.array([-1, 0, -1])
    arr2 = 0.02 * np.array([-1, 0, 1])

    arr1[::2] += np.random.normal(0, 0.005, 2)
    arr2[::2] += np.random.normal(0, 0.005, 2)

    x, y = xaxis.get_data()
    if '+' in str(xaxis_arrow):
        ax.plot(x[-1] + arr1 * xspan * aspect,
                y[-1] + arr2 * yspan,
                color='k', lw=2)
    if '-' in str(xaxis_arrow):
        ax.plot(x[0] - arr1 * xspan * aspect,
                y[0] - arr2 * yspan,
                color='k', lw=2)

    x, y = yaxis.get_data()
    if '+' in str(yaxis_arrow):
        ax.plot(x[-1] + arr2 * xspan * aspect,
                y[-1] + arr1 * yspan,
                color='k', lw=2)
    if '-' in str(yaxis_arrow):
        ax.plot(x[0] - arr2 * xspan * aspect,
                y[0] - arr1 * yspan,
                color='k', lw=2)

    # Change all the fonts to humor-sans.
    prop = fm.FontProperties(fname='Humor-Sans.ttf', size=16)
    for text in ax.texts:
        text.set_fontproperties(prop)
    
    # modify legend
    leg = ax.get_legend()
    if leg is not None:
        leg.set_frame_on(False)
        
        for child in leg.get_children():
            if isinstance(child, pl.Line2D):
                x, y = child.get_data()
                child.set_data(xkcd_line(x, y, mag=10, f1=100, f2=0.001))
                child.set_linewidth(2 * child.get_linewidth())
            if isinstance(child, pl.Text):
                child.set_fontproperties(prop)
    
    # Set the axis limits
    ax.set_xlim(xax_lim[0] - 0.1 * xspan,
                xax_lim[1] + 0.1 * xspan)
    ax.set_ylim(yax_lim[0] - 0.1 * yspan,
                yax_lim[1] + 0.1 * yspan)

    # adjust the axes
    ax.set_xticks([])
    ax.set_yticks([])      

    if expand_axes:
        ax.figure.set_facecolor(bgcolor)
        ax.set_axis_off()
        ax.set_position([0, 0, 1, 1])
    
    return ax

    
# def DaVinci_patch(patch):
    # cp = []
    # com = []
    
    # for i in xrange(5):
        # cp.extend([(x0, y0), (x1, y0), (x1, y1), (x0, y1),
          # (x0, y0), (x0, y0)])

        # com.extend([Path.LINETO,
           # Path.LINETO,
           # Path.LINETO,
           # Path.LINETO,
           # Path.LINETO,
           # Path.LINETO])
           
    # x,y = xkcd_line(cp[0][:], cp[1][:])

    # cp = [ (x[i], y[i]) for i in xrange(len(x))]
    
    # com[0] = Path.MOVETO
    # com[1] = Path.CLOSEPOLY
           
    # path = Path(cp, com)

    # return path
from matplotlib.image import BboxImage
class PathClippedImagePatch(patches.PathPatch):
    """
    The given image is used to draw the face of the patch. Internally,
    it uses BboxImage whose clippath set to the path of the patch.

    FIXME : The result is currently dpi dependent.
    """
    def __init__(self, path, bbox_image, **kwargs):
        patches.PathPatch.__init__(self, path, **kwargs)
        self._init_bbox_image(bbox_image)

    def set_facecolor(self, color):
        """simply ignore facecolor"""
        patches.PathPatch.set_facecolor(self, "none")

    def _init_bbox_image(self, im):

        bbox_image = BboxImage(self.get_window_extent,
                               norm = None,
                               origin=None,
                               )
        bbox_image.set_transform(transforms.IdentityTransform())

        bbox_image.set_data(im)
        self.bbox_image = bbox_image

    def draw(self, renderer=None):


        # the clip path must be updated every draw. any solution? -JJ
        self.bbox_image.set_clip_path(self._path, self.get_transform())
        self.bbox_image.draw(renderer)

        patches.PathPatch.draw(self, renderer)
        
def DaVincify(ax, mag=1.0,
            f1=50, f2=0.01, f3=15,
            bgcolor='#F2EDDC',
            xaxis_loc=None,
            yaxis_loc=None,
            xaxis_arrow='+',
            yaxis_arrow='+',
            ax_extend=0.1,
            expand_axes=False):
    """Make axis look hand-drawn

    This adjusts all lines, text, legends, and axes in the figure to look
    like xkcd plots.  Other plot elements are not modified.
    
    Parameters
    ----------
    ax : Axes instance
        the axes to be modified.
    mag : float
        the magnitude of the distortion
    f1, f2, f3 : int, float, int
        filtering parameters.  f1 gives the size of the window, f2 gives
        the high-frequency cutoff, f3 gives the size of the filter
    xaxis_loc, yaxis_log : float
        The locations to draw the x and y axes.  If not specified, they
        will be drawn from the bottom left of the plot
    xaxis_arrow, yaxis_arrow : str
        where to draw arrows on the x/y axes.  Options are '+', '-', '+-', or ''
    ax_extend : float
        How far (fractionally) to extend the drawn axes beyond the original
        axes limits
    expand_axes : bool
        if True, then expand axes to fill the figure (useful if there is only
        a single axes in the figure)
    """
    # Get axes aspect
    ext = ax.get_window_extent().extents
    aspect = (ext[3] - ext[1]) / (ext[2] - ext[0])

    xlim = ax.get_xlim()
    ylim = ax.get_ylim()

    xspan = xlim[1] - xlim[0]
    yspan = ylim[1] - xlim[0]

    xax_lim = (xlim[0] - ax_extend * xspan,
               xlim[1] + ax_extend * xspan)
    yax_lim = (ylim[0] - ax_extend * yspan,
               ylim[1] + ax_extend * yspan)

    if xaxis_loc is None:
        xaxis_loc = ylim[0]

    if yaxis_loc is None:
        yaxis_loc = xlim[0]

    # Draw axes
    xaxis = pl.Line2D([xax_lim[0], xax_lim[1]], [xaxis_loc, xaxis_loc],
                      linestyle='-', color="#362A1C", linewidth = 0.7)
    yaxis = pl.Line2D([yaxis_loc, yaxis_loc], [yax_lim[0], yax_lim[1]],
                      linestyle='-', color="#362A1C", linewidth = 0.7)

    # Label axes3, 0.5, 'hello', fontsize=14)
    ax.text(xax_lim[1], xaxis_loc - 0.05 * yspan, ax.get_xlabel(),
            fontsize=14, ha='right', va='top')
    ax.text(yaxis_loc - 0.05 * xspan, yax_lim[1], ax.get_ylabel(),
            fontsize=14, ha='right', va='top')
    ax.set_xlabel('')
    ax.set_ylabel('')

    # Add title
    ax.text(0.5 * (xax_lim[1] + xax_lim[0]), yax_lim[1]*1.08,
            ax.get_title(),
            ha='center', va='bottom', fontsize=26)
    ax.set_title('')

    grids = []
    for yt in ax.get_yticks():
        if yt>= ylim[0] and yt<=ylim[1]:
            grids.append(pl.Line2D([xlim[0]-0.05 * random() * xspan, xlim[1]]+0.05 * random() * xspan, [yt, yt],
                          linestyle='-', color="#111111", linewidth = 0.1))
    
    for xt in ax.get_xticks():
        if xt>= xlim[0] and xt<=xlim[1]:
            grids.append(pl.Line2D([xt, xt], [ylim[0]-0.05 * random() *yspan, ylim[1]+0.05 * random() * yspan],
                          linestyle='-', color="#111111", linewidth = 0.1))

    # Draw arrow-heads at the end of axes lines
    arr1 = 0.03 * np.array([-1, 0, -1])
    arr2 = 0.02 * np.array([-1, 0, 1])

    #arr1[::2] += np.random.normal(0, 0.001, 2)
    #arr2[::2] += np.random.normal(0, 0.001, 2)

    x, y = xaxis.get_data()
    if '+' in str(xaxis_arrow):
        ax.plot(x[-1] + arr1 * xspan * aspect,
                y[-1] + arr2 * yspan,
                color="#362A1C", lw=1)
    if '-' in str(xaxis_arrow):
        ax.plot(x[0] - arr1 * xspan * aspect,
                y[0] - arr2 * yspan,
                color="#362A1C", lw=1)

    x, y = yaxis.get_data()
    if '+' in str(yaxis_arrow):
        ax.plot(x[-1] + arr2 * xspan * aspect,
                y[-1] + arr1 * yspan,
                color="#362A1C", lw=1)
    if '-' in str(yaxis_arrow):
        ax.plot(x[0] - arr2 * xspan * aspect,
                y[0] - arr1 * yspan,
                color="#362A1C", lw=1)

    Nlines = len(ax.lines)
    lines = grids+[xaxis, yaxis] + [ax.lines.pop(0) for i in range(Nlines)]
            
    fg = []
    shade = []
    
    for line in lines:
        l = DaVinci_line(line, "b",  xlim = xlim, ylim = ylim)
        fg.extend(l[0])
        shade.extend(l[1])
    
    for l in shade:
        ax.add_line(l)
        
    for l in fg:
        ax.add_line(l)
                
    # Change all the fonts to humor-sans.
    

    arr = np.arange(256).reshape(1,256)/256
    
    prop = fm.FontProperties(fname='davinci.ttf', size=34)
    for text in ax.texts:
        text.set_fontproperties(prop)
        text.set_path_effects([ patheffects.withStroke(linewidth=5,
                                foreground=bgcolor)])
    
    # modify legend
    leg = ax.get_legend()
    if leg is not None:
        #leg.set_frame_on(False)
        #leg.set_zorder(1)
        leg.draggable(True)
        frame = leg.get_frame()
        frame.set_boxstyle('round')
        frame.set_lw(0)
        frame.set_facecolor(bgcolor)
        
        #corners = leg.get_bbox_to_anchor().corners()
        #corners = [corners[0], corners[1],corners[3],corners[2]]
        #transform = transforms.BboxTransformFrom(leg.get_bbox_to_anchor())
        # leg_ax = frame.get_axes()
        #leg.set_frame_on(False)
        
        # com = [Path.MOVETO,
           # Path.LINETO,
           # Path.LINETO,
           # Path.CLOSEPOLY]
        
        #g = patches.Polygon(corners)
        # print dir(frame)
        
        
        #g.set_clip_box(None)
        # print dir(leg.get_bbox_to_anchor())
        #g.set_transform(transforms.IdentityTransform())
        
        pl.draw()
        
        bbox = leg.legendPatch.get_bbox().inverse_transformed(ax.transAxes)
        g = patches.Rectangle((bbox.x0, bbox.y0), 
                      bbox.width*1.2, bbox.height*1.2, 
                      fc='red', transform=ax.transAxes, zorder=50)
        
        ax.add_patch(g)
        
        print getp(g)
        print getp(frame)
        print getp(leg)
        
        for child in leg.get_children():
            if isinstance(child, pl.Line2D):
                x, y = child.get_data()
                child.set_data(xkcd_line(x, y, mag=10, f1=100, f2=0.001))
                child.set_linewidth(2 * child.get_linewidth())
            if isinstance(child, pl.Text):
                child.set_fontproperties(prop)
    
    # Set the axis limits
    ax.set_xlim(xax_lim[0] - 0.1 * xspan,
                xax_lim[1] + 0.1 * xspan)
    ax.set_ylim(yax_lim[0] - 0.1 * yspan,
                yax_lim[1] + 0.1 * yspan)

    # adjust the axes
    #ax.set_xticks([])
    #ax.set_yticks([])      

    if expand_axes:
        ax.figure.set_facecolor(bgcolor)
        ax.set_axis_off()
        ax.set_position([0.03, 0.03, 0.92, 0.92])
    
    return ax
    
    
# np.random.seed(0)

# ax = pl.axes()

# x = np.linspace(-0.01, 1.1, 50)
# def cannon(x, angle):
    # return x / np.tan(angle) - 0.5 * x**2/(np.sin(angle)**2)
    
# ax.plot(x, cannon(x,0.3)*2, lw=1)
# ax.plot(x, cannon(x,0.45)*2, lw=1)
# ax.plot(x, cannon(x,0.6)*2, lw=1)
# ax.plot(x, cannon(x,0.75)*2,'k', lw=1)
# ax.plot(x, cannon(x,0.9)*2,'b', lw=1)

# #ax.text(0.9, 0.4, "embarrassment")

# ax.set_title('Study of the trajectory of a cannon ball')
# ax.set_xlabel('Distance')
# ax.set_ylabel('Height')

# ax.legend(loc='top right', fancybox = True)

# ax.set_xlim(-0.05, 1.2)
# ax.set_ylim(-.1, 1.0)
# pl.draw()

# #XKCDify the axes -- this operates in-place
# DaVincify(ax, xaxis_loc=0.0, yaxis_loc=0.0,
        # xaxis_arrow='+', yaxis_arrow='+',
        # expand_axes=True)
# pl.show()